import pandas as pd
import joblib
import os

# 加载模型
try:
    model = joblib.load('models/best_imbalanced_model.pkl')
    print("模型加载成功")
except FileNotFoundError:
    print("错误：未找到模型文件，请确保最佳模型已训练并保存在models目录下")
    exit(1)
except Exception as e:
    print(f"错误：加载模型时发生未知错误 - {e}")
    exit(1)

# 加载特征数据，用于提取特征列
try:
    features_data = pd.read_csv('final_features_with_pca.csv')
    print("特征数据加载成功")
except FileNotFoundError:
    print("错误：未找到特征数据文件 'final_features_with_pca.csv'")
    exit(1)
except Exception as e:
    print(f"错误：加载特征数据时发生未知错误 - {e}")
    exit(1)

# 获取特征列名（排除user_id和seller_id）
feature_columns = features_data.columns.tolist()
feature_columns = [col for col in feature_columns if col not in ['user_id', 'seller_id']]

# 加载待预测数据
try:
    data = pd.read_csv('test_without_label.csv')
    print(f"待预测数据加载成功，形状: {data.shape}")
except FileNotFoundError:
    print("错误：未找到待预测数据文件 'test_without_label.csv'")
    exit(1)
except Exception as e:
    print(f"错误：加载待预测数据时发生未知错误 - {e}")
    exit(1)

# 检查数据列是否符合预期
required_columns = ['user_id', 'merchant_id']
for col in required_columns:
    if col not in data.columns:
        print(f"错误：待预测数据中缺少必要的列 '{col}'")
        exit(1)

# 重命名列以匹配模型训练时的列名
data = data.rename(columns={'merchant_id': 'seller_id'})

# 合并特征数据
try:
    merged_data = pd.merge(data, features_data, on=['user_id', 'seller_id'], how='left')
    print(f"数据合并完成，形状: {merged_data.shape}")
except Exception as e:
    print(f"错误：合并数据时发生未知错误 - {e}")
    exit(1)

# 提取特征
X = merged_data[feature_columns]

# 检查是否有缺失值
missing_values = X.isnull().sum().sum()
if missing_values > 0:
    print(f"警告：特征数据中包含{missing_values}个缺失值，将使用0填充")
    X = X.fillna(0)

# 预测
try:
    print("开始预测...")
    # 获取标签为1的概率
    probs = model.predict_proba(X)[:, 1]
    # 获取预测标签（默认阈值0.5）
    labels = model.predict(X)

    print("预测完成")
except Exception as e:
    print(f"错误：预测时发生未知错误 - {e}")
    exit(1)

# 添加预测结果到原始数据
data['prob'] = probs
data['label'] = labels

# 保存结果
output_file = 'train_with_predictions.csv'
try:
    data.to_csv(output_file, index=False)
    print(f"预测结果已保存到 {output_file}")
    print(f"结果文件包含 {len(data)} 条记录")
except Exception as e:
    print(f"错误：保存预测结果时发生未知错误 - {e}")